from model.my_chat_model import ChatModel

def rewrite_query(original_question:str)->str:
    """
    使用Qwen-Max 对用户问题进行语义补全与意图澄清
    :param original_question: 原始用户问题
    :return:
    """
    chat = ChatModel()
    llm = chat.get_line_model()

    prompt = f"""
    你是一个专业的查询重写助手。请分析用户的原始问题，明确其真实意图，并将其重写为一个更清晰、包含关键词、适合知识库检索和工具调用的查询语句。
    
    要求：
    1. 保留原意，不得曲解；
    2. 补充可能的关键词（如：企业、股东、高管、投资关系、持股关系、任职关系、关联关系、企业背景等）；
    3. 输出必须是一句完整的、可执行的查询语句，不要解释；
    4. 如果涉及时间，请明确是否需要当前时间。

    示例：
    输入：公司名称为公司4的公司怎么样？
    输出：请查询公司名称为公司4的公司所经营的行业、风险水平高低、目前的状态、企业背景信息。
    
    输入：高管名称为高管4的高管怎么样？
    输出：请查询高管名称为高管4的高管的出生日期，性别，在哪些公司任职，以及胜任的职位。
    
    输入：股东名称为自然人股东4的股东怎么样？
    输出：股东名称为自然人股东4的股东的国籍，持有着哪些公司的股票，对各个公司股票的持股比例是多少，以及他和哪些股东有关系，是什么关系。
    
    输入：现在几点？
    输出：请获取当前时刻的时间。

    现在请处理以下问题：
    输入：{original_question}
    输出：
    """

    #调用大模型生成重写后的问题
    response = llm.invoke(prompt)
    rewritten_query = response.content.strip()

    #防止空值或无效输出
    if not rewritten_query or len(rewritten_query) < 5:
        return original_question
    return rewritten_query